import os
from pathlib import Path

def parse_filename(filename):
    """解析CCPD文件名，提取车牌位置（忽略颜色）"""
    parts = filename.split('-')
    if len(parts) < 6:
        raise ValueError(f"Invalid filename format: {filename}")

    # 解析车牌四个角点坐标
    corners_part = parts[2].split('_')
    corners = []
    for corner in corners_part[:4]:
        x, y = map(int, corner.split('&'))
        corners.append((x, y))

    return {
        'corners': corners,
        'plate_number': parts[5]  # 保留车牌号信息（可选）
    }

def corners_to_yolo(corners, img_width, img_height):
    """将角点转换为YOLO格式的矩形框"""
    x_coords = [x for x, y in corners]
    y_coords = [y for x, y in corners]

    x_min, x_max = min(x_coords), max(x_coords)
    y_min, y_max = min(y_coords), max(y_coords)

    center_x = (x_min + x_max) / 2 / img_width
    center_y = (y_min + y_max) / 2 / img_height
    width = (x_max - x_min) / img_width
    height = (y_max - y_min) / img_height

    return center_x, center_y, width, height

def process_subset(subset_dir, output_root, img_width=720, img_height=1160):
    """
    处理单个子集（train/val/test）
    :param subset_dir: 子集目录路径
    :param output_root: 输出目录
    """
    subset_name = os.path.basename(subset_dir)
    output_dir = os.path.join(output_root, subset_name, 'labels')
    os.makedirs(output_dir, exist_ok=True)

    processed_count = 0
    for img_file in Path(subset_dir).glob('*.jpg'):
        try:
            # 解析文件名并生成YOLO标注（固定class_id=0）
            info = parse_filename(img_file.stem)
            cx, cy, w, h = corners_to_yolo(info['corners'], img_width, img_height)

            # 写入标注文件（所有车牌类别设为0）
            txt_file = os.path.join(output_dir, f"{img_file.stem}.txt")
            with open(txt_file, 'w') as f:
                f.write(f"0 {cx:.6f} {cy:.6f} {w:.6f} {h:.6f}")

            processed_count += 1
        except Exception as e:
            print(f"Error processing {img_file.name}: {str(e)}")

    print(f"Processed {processed_count} files in {subset_name}")

def prepare_yolo_dataset(dataset_root, output_root):
    """
    处理完整的CCPD数据集（所有车牌标注为类别0）
    :param dataset_root: 数据集根目录
    :param output_root: 输出目录
    """
    os.makedirs(output_root, exist_ok=True)

    for subset in ['train', 'val', 'test']:
        subset_dir = os.path.join(dataset_root, subset)
        if os.path.exists(subset_dir):
            print(f"\nProcessing {subset} set...")
            process_subset(subset_dir, output_root)
        else:
            print(f"Warning: {subset} directory not found, skipping")

if __name__ == "__main__":
    # 配置路径（根据实际情况修改）
    DATASET_ROOT = r"D:\license_plate\datasets\ccpd_blue"  # 数据集路径
    OUTPUT_ROOT = r"D:\license_plate\datasets\ccpd_blue\labels_blue"  # 输出路径

    # 执行转换
    prepare_yolo_dataset(DATASET_ROOT, OUTPUT_ROOT)
    print("\nConversion completed!")
    print(f"YOLO format dataset saved to: {OUTPUT_ROOT}")